import sys
import cv2
from PySide6.QtCore import Qt, QTimer, Signal, QObject
from PySide6.QtGui import QImage, QPixmap
from PySide6.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget, QMessageBox
from qfluentwidgets import PushButton


class VideoStream(QObject):
    new_frame = Signal(QImage)
    error_occurred = Signal(str)

    def __init__(self, rtsp_url):
        super().__init__()
        self.stream_url = rtsp_url  # 直接使用RTSP URL
        self.cap = None
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.retry_count = 0
        self.max_retry = 5

        self.open_stream()

    def open_stream(self):
        """尝试打开RTSP视频流"""
        if self.cap:
            self.cap.release()

        # 创建VideoCapture对象并设置RTSP参数
        self.cap = cv2.VideoCapture(self.stream_url)

        # 关键设置：使用TCP传输（避免UDP丢包问题），设置超时和缓冲区
        # self.cap.set(cv2.CAP_PROP_RTSP_TRANSPORT, 1)  # 1 = TCP传输模式
        self.cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 5000)  # 5秒超时
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 最小化缓冲区

        # 打开RTSP流
        self.cap.open(self.stream_url)

        if not self.cap.isOpened():
            self.retry_count += 1
            if self.retry_count <= self.max_retry:
                QTimer.singleShot(2000, self.open_stream)  # 2秒后重试
            else:
                self.error_occurred.emit(f"无法打开RTSP视频流: {self.stream_url}")
        else:
            self.retry_count = 0
            self.timer.start(30)  # 约30fps

    def update_frame(self):
        if not self.cap or not self.cap.isOpened():
            return

        ret, frame = self.cap.read()
        if ret:
            # 转换为黑白图像
            # gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            # h, w = gray_image.shape
            # qt_image = QImage(gray_image.data, w, h, w, QImage.Format.Format_Grayscale8)

            # 转换为RGB图像
            rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            qt_image = QImage(rgb_image.data, w, h, w * ch, QImage.Format.Format_RGB888)

            self.new_frame.emit(qt_image)
        else:
            # 读取失败时重新连接
            self.open_stream()

    def stop(self):
        self.timer.stop()
        if self.cap:
            self.cap.release()


class CameraInterface(QWidget):
    def __init__(self, rtsp_url="rtsp://127.0.0.1:8554/stream"):
        super().__init__()
        self.setObjectName("camera_interface")
        self.setWindowTitle("RTSP视频流接收器")
        self.resize(800, 600)

        # 创建UI
        self.label = QLabel("等待RTSP视频流...")
        self.label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.label.setStyleSheet("font-size: 20px; color: #888;")

        # 创建按钮
        self.listen_button = PushButton("开始监听RTSP")
        self.listen_button.clicked.connect(lambda: self.start_listen(rtsp_url))

        layout = QVBoxLayout()
        layout.setContentsMargins(20, 50, 20, 20)
        layout.addWidget(self.label)
        layout.addWidget(self.listen_button)
        self.setLayout(layout)

    def start_listen(self, rtsp_url):
        self.listen_button.setEnabled(False)
        self.label.setText(f"正在连接: {rtsp_url}...")
        # 创建视频流
        self.stream = VideoStream(rtsp_url)
        self.stream.new_frame.connect(self.display_frame)
        self.stream.error_occurred.connect(self.display_error)

    def display_frame(self, image):
        pixmap = QPixmap.fromImage(image)
        self.label.setPixmap(
            pixmap.scaled(
                self.label.width(),
                self.label.height(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation,
            )
        )

    def display_error(self, message):
        QMessageBox.critical(self, "RTSP视频流错误", message)
        self.label.setText(f"错误: {message}")
        self.listen_button.setEnabled(True)  # 出错时重新启用按钮

    def closeEvent(self, event):
        if hasattr(self, "stream") and self.stream:
            self.stream.stop()
        super().closeEvent(event)


if __name__ == "__main__":
    app = QApplication(sys.argv)

    # 使用命令行参数指定RTSP URL
    rtsp_url = "rtsp://127.0.0.1:8554/stream"  # 默认URL
    if len(sys.argv) > 1:
        rtsp_url = sys.argv[1]

    player = CameraInterface(rtsp_url)
    player.show()
    sys.exit(app.exec())
